# frozen_string_literal: true

module EE
  module Gitlab
    module BackgroundMigration
      module BackfillDismissalReasonInVulnerabilityReads
        extend ActiveSupport::Concern
        extend ::Gitlab::Utils::Override

        prepended do
          operation_name :backfill_dismissal_reason
          scope_to ->(relation) { relation.where(state: 2, dismissal_reason: nil) }
          feature_category :vulnerability_management
        end

        class VulnerabilitiesStateTransition < ::ApplicationRecord # rubocop:disable Style/Documentation
          self.table_name = 'vulnerability_state_transitions'
        end

        class VulnerabilitiesRead < ::ApplicationRecord # rubocop:disable Style/Documentation
          self.table_name = 'vulnerability_reads'
        end

        class Vulnerability < ::ApplicationRecord # rubocop:disable Style/Documentation
          self.table_name = 'vulnerabilities'
        end

        DETECTED_STATE = 1
        DISMISSED_STATE = 2

        # So we might need to dwo two passes here
        # 1. If there's a dismissing Vulnerabilities::StateTransition then we use it in
        #     `create_attributes_from_state_transitions`
        #
        # 2. If not, then we fetch the Vulnerability and use `dismissed_at` and `dismissed_by_id` to
        #     create an appropriate Vulnerabilities::StateTransition records, since we don't have
        #     a `dismissal_reason` information at this point then we don't need to update the
        #     Vulnerabilities::Read record
        override :perform
        def perform
          each_sub_batch do |sub_batch|
            vulnerability_read_attributes, vulnerability_ids = prepare_upsert_data(sub_batch)

            VulnerabilitiesRead.upsert_all(vulnerability_read_attributes) if vulnerability_read_attributes.any?

            next unless vulnerability_ids.any?

            vulnerability_state_transition_attributes = prepare_state_transitions_from_vulnerabilities(
              vulnerability_ids
            )

            VulnerabilitiesStateTransition.upsert_all(
              vulnerability_state_transition_attributes
            )
          end
        end

        private

        def prepare_upsert_data(sub_batch)
          dismissed_reads, state_transitions = load_data(sub_batch)

          prepare_attributes_from_state_transitions(dismissed_reads, state_transitions)
        end

        def load_data(sub_batch)
          dismissed_reads = VulnerabilitiesRead.where(id: sub_batch)
          vulnerability_ids = dismissed_reads.pluck(:vulnerability_id)
          state_transitions = VulnerabilitiesStateTransition
            .where(vulnerability_id: vulnerability_ids, to_state: DISMISSED_STATE)
            .order(id: :desc)

          [dismissed_reads, state_transitions]
        end

        def prepare_attributes_from_state_transitions(dismissed_reads, state_transitions)
          attributes = dismissed_reads.map do |read|
            dismissing_state_transition = state_transitions.find do |st|
              st.vulnerability_id == read.vulnerability_id
            end

            if dismissing_state_transition
              read.attributes.merge(dismissal_reason: dismissing_state_transition.dismissal_reason)
            else
              read.vulnerability_id
            end
          end

          attributes.partition { |e| e.is_a?(Hash) }
        end

        def prepare_state_transitions_from_vulnerabilities(ids)
          vulnerabilities = Vulnerability.where(id: ids)
          vulnerabilities.map do |vulnerability|
            {
              vulnerability_id: vulnerability.id,
              from_state: DETECTED_STATE,
              to_state: DISMISSED_STATE,
              created_at: vulnerability.dismissed_at,
              updated_at: Time.current,
              author_id: vulnerability.dismissed_by_id
            }
          end
        end
      end
    end
  end
end
